import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from python_ai.common.xcommon import *
from python_ai.CV_2.app.teachers_day17.my.inception_keras_oop_cifar10 import InceptionNet
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, optimizers, losses
from sklearn.model_selection import train_test_split
import os

np.random.seed(3)
FILE_NAME = os.path.basename(__file__)

vcode_dir_train = '../../../../large_data/DL1/_many_files/vcode_data/train'
vcode_dir_test = '../../../../large_data/DL1/_many_files/vcode_data/test'

def load_dir(vcode_dir):
    names = os.listdir(vcode_dir)
    n_names = len(names)
    names = np.array(names)
    idx = np.random.permutation(n_names)  # for shuffle
    names = names[idx]
    x = []
    y = []
    for name in names:
        digits, ext = os.path.splitext(name)
        if ext.lower() != '.jpg':
            continue
        digits_list = list(digits)
        if len(digits_list) != 4:
            continue
        digits_arr = []
        try:
            for di in digits_list:
                digits_arr.append(int(di))
        except ValueError as ex:
            continue

        y.append(digits_arr)

        path = os.path.join(vcode_dir, name)
        img = cv.imread(path, cv.IMREAD_GRAYSCALE)
        x.append(img)

    x = np.float32(x) / 255.
    y = np.uint8(y)
    x = np.expand_dims(x, 3)
    y = keras.utils.to_categorical(y, 10)
    y = y.reshape((-1, 40))
    return x, y

x_train, y_train = load_dir(vcode_dir_train)
x_test, y_test = load_dir(vcode_dir_test)

print('x_train', x_train.shape)
print('y_train', y_train.shape)
print('x_test', x_test.shape)
print('y_test', y_test.shape)

model = keras.Sequential([
    layers.Conv2D(6, (5, 5)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(16, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(32, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Conv2D(64, (3, 3)),
    layers.BatchNormalization(),
    layers.ReLU(),
    layers.MaxPool2D(strides=(2, 2), padding='same'),

    layers.Flatten(),
    layers.Dense(120, activation=activations.relu),
    layers.Dense(84, activation=activations.relu),
    layers.Dense(40, activation='sigmoid'),
])
model.build(input_shape=(None, 60, 160, 1))
model.summary(line_length=120)


eps = 1e-20
def my_loss(y_true, y_pred):
    mat = y_true * tf.math.log(y_pred + eps) + (1 - y_true) * tf.math.log(1 - y_pred + eps)
    col = tf.reduce_sum(mat, axis=1)
    j = - tf.reduce_mean(col)
    return j


def my_acc(y_true, y_pred):
    y_true = tf.reshape(y_true, (-1, 4, 10))
    y_pred = tf.reshape(y_pred, (-1, 4, 10))
    y_true = tf.argmax(y_true, axis=2)
    y_pred = tf.argmax(y_pred, axis=2)
    eq = y_true == y_pred
    eq = tf.reduce_all(eq, axis=1)
    acc = tf.reduce_mean(tf.cast(eq, dtype=tf.float32))
    return acc


# y_pred = model.predict(x_test)
# # y_pred = model(x_test)  # tensorflow.python.framework.errors_impl.UnknownError: Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above. [Op:Conv2D]
# j = my_loss(y_test, y_pred)
# print(j)
# acc = my_acc(y_test, y_pred)
# print(acc)
# sys.exit(0)

model.compile(loss=my_loss,
              optimizer=optimizers.Adam(lr=0.00001),
              metrics=my_acc)

history=model.fit(
    x_train, y_train, batch_size=64, epochs=200,
    validation_data=(x_test, y_test), validation_batch_size=64,
    callbacks=[keras.callbacks.TensorBoard(log_dir='_log/' + FILE_NAME, update_freq='batch', profile_batch=0)],

)

score=model.evaluate(x_test, y_test, batch_size=64)
print('accuracy',score[1])
print('loss',score[0])
